import torch

# 定义数据集
x_data = torch.tensor([[1.0], [2.0], [3.0]], requires_grad=False)
y_target = torch.tensor([[2.0], [10.0], [15.0]], requires_grad=False)

# 定义模型参数
w = torch.randn(1, requires_grad=True)  # 权重
b = torch.randn(1, requires_grad=True)  # 偏置

# 定义学习率和迭代次数
learning_rate = 0.01
epochs = 500

for epoch in range(epochs):
    # 正向传播：计算预测值 y = wx + b
    y_pred = w * x_data + b

    # 计算损失 (均方误差)
    loss = torch.mean((y_pred - y_target) ** 2)

    # 反向传播：计算梯度
    loss.backward()

    # 更新参数
    with torch.no_grad():
        print(w,w.grad)
        w -= learning_rate * w.grad
        b -= learning_rate * b.grad

        # 手动清零梯度
        w.grad.zero_()
        b.grad.zero_()

    if (epoch+1) % 50 == 0:
        print(f'Epoch [{epoch+1}/{epochs}], Loss: {loss.item():.4f}')

print(f'Trained: w = {w.item()}, b = {b.item()}')

# 测试模型
test_x = torch.tensor([[4.0]])
predicted_y = w * test_x + b
print(f'Predicted y for x=4: {predicted_y.item()}')